import os
import torch
import torch.nn as nn
import torchvision.transforms as transforms
from torchvision import datasets, models
from torch.utils.data import DataLoader
import matplotlib.pyplot as plt

# 创建保存模型的文件夹
save_dir = './models'
os.makedirs(save_dir, exist_ok=True)  # 确保文件夹存在

# 加载预训练的 ConvNeXt-Tiny 模型
model = models.convnext_tiny(weights=None)

# 数据预处理和增强
transform_train = transforms.Compose([
    transforms.RandomResizedCrop((224, 224)),  # 随机裁剪并调整大小到 224x224
    transforms.RandomHorizontalFlip(),  # 随机水平翻转
    transforms.ToTensor(),  # 转换为张量
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),  # 归一化
])

transform_val = transforms.Compose([
    transforms.Resize((224, 224)),  # 调整大小到 224x224
    transforms.ToTensor(),  # 转换为张量
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),  # 归一化
])

data_dir = '/root/dr/CI/data_tiny_imagenet/tiny-imagenet-200/'
# 加载训练集
train_dataset = datasets.ImageFolder(root=data_dir + 'train/', transform=transform_train)
train_loader = DataLoader(train_dataset, batch_size=128, shuffle=True, num_workers=4)

# 加载验证集
val_dataset = datasets.ImageFolder(root=data_dir + 'val/', transform=transform_val)
val_loader = DataLoader(val_dataset, batch_size=128, shuffle=False, num_workers=4)

# 替换最后的分类层以适应 tiny_imagenet 的 200 个类别，添加 Dropout
# ConvNeXt-Tiny 的最后一层输出是 768 维
model.classifier[2] = nn.Sequential(
    nn.Dropout(0.5),  # 加入 Dropout，防止过拟合
    nn.Linear(in_features=768, out_features=200)
)

# 将模型移到 GPU（如果可用）
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = model.to(device)
print(f'Training on device: {device}')

# 定义损失函数和优化器，加入 L2 正则（weight_decay）
criterion = nn.CrossEntropyLoss()
# 将初始学习率设为 0.001，并使用 L2 正则化
optimizer = torch.optim.Adam(model.parameters(), lr=0.001, weight_decay=1e-4)

# 使用余弦退火学习率调度器
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=200)

# 记录训练过程中的损失、准确率和学习率
train_losses = []
train_top1_accuracies = []
train_top5_accuracies = []
val_top1_accuracies = []
val_top5_accuracies = []
learning_rates = []  # 记录学习率

# 定义计算 Top-1 和 Top-5 准确率的函数
def calculate_top1_top5_accuracy(loader, model):
    top1_correct = 0
    top5_correct = 0
    total = 0
    model.eval()  # 切换到评估模式
    with torch.no_grad():
        for inputs, labels in loader:
            inputs, labels = inputs.to(device), labels.to(device)
            outputs = model(inputs)
            
            # 计算 Top-1 准确率
            _, top1_pred = torch.max(outputs, 1)
            top1_correct += (top1_pred == labels).sum().item()
            
            # 计算 Top-5 准确率
            _, top5_pred = torch.topk(outputs, 5, dim=1)
            top5_correct += (top5_pred == labels.view(-1, 1)).sum().item()

            total += labels.size(0)
    
    top1_accuracy = top1_correct / total
    top5_accuracy = top5_correct / total
    return top1_accuracy, top5_accuracy

num_epochs = 300  #

# 训练模型
for epoch in range(num_epochs):
    model.train()  # 切换到训练模式
    running_loss = 0.0
    top1_correct = 0
    top5_correct = 0
    total = 0
    for inputs, labels in train_loader:
        inputs, labels = inputs.to(device), labels.to(device)

        # 前向传播
        optimizer.zero_grad()
        outputs = model(inputs)
        loss = criterion(outputs, labels)
        
        # 反向传播和优化
        loss.backward()
        optimizer.step()

        running_loss += loss.item()

        # 计算训练 Top-1 准确率
        _, top1_pred = torch.max(outputs, 1)
        top1_correct += (top1_pred == labels).sum().item()

        # 计算训练 Top-5 凼确率
        _, top5_pred = torch.topk(outputs, 5, dim=1)
        top5_correct += (top5_pred == labels.view(-1, 1)).sum().item()

        total += labels.size(0)

    # 更新学习率调度器
    scheduler.step()

    # 获取当前学习率并记录
    current_lr = scheduler.get_last_lr()[0]  # 取第一个学习率（如果有多个参数组）
    learning_rates.append(current_lr)

    train_loss = running_loss / len(train_loader)
    train_top1_acc = top1_correct / total
    train_top5_acc = top5_correct / total
    val_top1_acc, val_top5_acc = calculate_top1_top5_accuracy(val_loader, model)

    train_losses.append(train_loss)
    train_top1_accuracies.append(train_top1_acc)
    train_top5_accuracies.append(train_top5_acc)
    val_top1_accuracies.append(val_top1_acc)
    val_top5_accuracies.append(val_top5_acc)

    print(f"Epoch {epoch+1}/{num_epochs}, Train Loss: {train_loss:.4f}, "
          f"Train Top-1 Acc: {train_top1_acc:.4f}, Train Top-5 Acc: {train_top5_acc:.4f}, "
          f"Val Top-1 Acc: {val_top1_acc:.4f}, Val Top-5 Acc: {val_top5_acc:.4f}, "
          f"LR: {current_lr:.6f}")

    # 每50个epoch保存一次模型
    if (epoch + 1) % 50 == 0:
        save_path = os.path.join(save_dir, f'convnext_tinyimagenet_epoch_{epoch+1}.pth')
        torch.save(model.state_dict(), save_path)
        print(f"Model saved at {save_path}")

# 绘制训练损失、训练精度、验证精度和学习率的曲线图
epochs_range = range(1, num_epochs + 1)
plt.figure(figsize=(12, 8))

# 绘制训练损失，并在同一张图上绘制学习率
ax1 = plt.subplot(2, 1, 1)
ax1.plot(epochs_range, train_losses, label='Train Loss', color='b')
ax1.set_xlabel('Epochs')
ax1.set_ylabel('Loss', color='b')
ax1.set_title('Train Loss and Learning Rate')
ax1.legend(loc='upper left')

# 创建 y 轴的双轴图，用于绘制学习率
ax2 = ax1.twinx()
ax2.plot(epochs_range, learning_rates, label='Learning Rate', color='r', linestyle=':')
ax2.set_ylabel('Learning Rate (Log Scale)', color='r')
ax2.set_yscale('log')  # 设置学习率为对数刻度
ax2.tick_params(axis='y', labelcolor='r')
ax2.legend(loc='upper right')

# 绘制训练和验证的 Top-1 和 Top-5 准确率在同一张图上
plt.subplot(2, 1, 2)
plt.plot(epochs_range, train_top1_accuracies, label='Train Top-1 Accuracy', linestyle='--')
plt.plot(epochs_range, train_top5_accuracies, label='Train Top-5 Accuracy', linestyle='--')
plt.plot(epochs_range, val_top1_accuracies, label='Val Top-1 Accuracy', linestyle='-')
plt.plot(epochs_range, val_top5_accuracies, label='Val Top-5 Accuracy', linestyle='-')
plt.xlabel('Epochs')
plt.ylabel('Accuracy')
plt.title('Top-1 and Top-5 Accuracy')
plt.legend()

# 保存图像到工作目录
plt.tight_layout()
plt.savefig('training_results_top1_top5_lr.png')
plt.show()

# 最终保存模型的 state_dict 到磁盘
final_model_path = os.path.join(save_dir, 'convnext_tinyimagenet_200epochs.pth')
torch.save(model.state_dict(), final_model_path)
print(f"Final model saved at {final_model_path}")